#include "myyolodetect.hpp"

static YOLO_V8 *yoloDetector = NULL;

cv::Rect cvRrect;

cv::Mat videoDetector(YOLO_V8*& p,cv::Mat &img);
int ReadCocoYaml(YOLO_V8*& p);

// Rect cvPointsToQRect(const cv::Point& topLeft, const cv::Point& bottomRight)
// {
//     // 创建QRect对象，使用cv::Point的x和y坐标
//     Rect rect(topLeft.x, topLeft.y, bottomRight.x - topLeft.x, bottomRight.y - topLeft.y);
//     return rect;
// }

// 实时视频检测
cv::Mat videoDetector(YOLO_V8 *&p, cv::Mat &img)
{

    std::vector<DL_RESULT> res;
    p->RunSession(img, res);

    for (auto &re : res)
    {
        cv::RNG rng(cv::getTickCount());
        cv::Scalar color(rng.uniform(0, 256), rng.uniform(0, 256), rng.uniform(0, 256));

        float confidence = floor(100 * re.confidence) / 100;
        if (confidence < 0.6)
        {
            continue;
        }
        cv::rectangle(img, re.box, color, 1);

        std::cout << std::fixed << std::setprecision(2);
        std::string label = p->classes[re.classId] + " " +
                            std::to_string(confidence).substr(0, std::to_string(confidence).size() - 4);
        // 标签绘制
        cv::rectangle(
            img,
            cv::Point(re.box.x, re.box.y - 25),
            cv::Point(re.box.x + label.length() * 15, re.box.y),
            color,
            cv::FILLED);

        cv::putText(
            img,
            label,
            cv::Point(re.box.x, re.box.y - 5),
            cv::FONT_HERSHEY_SIMPLEX,
            0.75,
            cv::Scalar(0, 0, 0),
            2);
        /* get QRcode rect */
        std::string name_id = p->classes[re.classId];
        if (re.classId == 2 && confidence >= 0.85)
        {
            cvRrect = cv::Rect(re.box);
        }
        //         QString labelstr=QString::fromStdString(std::to_string(re.classId)+p->classes[re.classId]);
        //         qDebug()<<"类别为："<<labelstr;
    }
    return img;
}

int ReadCocoYaml(YOLO_V8 *&p)
{
    // Open the YAML file
    std::ifstream file("coco.yaml");
    if (!file.is_open())
    {
        std::cerr << "Failed to open file" << std::endl;
        return 1;
    }

    // Read the file line by line
    std::string line;
    std::vector<std::string> lines;
    while (std::getline(file, line))
    {
        lines.push_back(line);
    }

    // Find the start and end of the names section
    std::size_t start = 0;
    std::size_t end = 0;
    for (std::size_t i = 0; i < lines.size(); i++)
    {
        if (lines[i].find("names:") != std::string::npos)
        {
            start = i + 1;
        }
        else if (start > 0 && lines[i].find(':') == std::string::npos)
        {
            end = i;
            break;
        }
    }

    // Extract the names
    std::vector<std::string> names;
    for (std::size_t i = start; i < end; i++)
    {
        std::stringstream ss(lines[i]);
        std::string name;
        std::getline(ss, name, ':'); // Extract the number before the delimiter
        std::getline(ss, name);      // Extract the string after the delimiter
        names.push_back(name);
    }

    p->classes = names;
    return 0;
}


void UnInit_onnx_yolov8(void)
{
    if (yoloDetector)
    {
        delete yoloDetector;
        yoloDetector = NULL;
    }
}

void Init_onnx_yolov8(void)
{
    yoloDetector = new YOLO_V8;
    ReadCocoYaml(yoloDetector);
    DL_INIT_PARAM params;
    params.rectConfidenceThreshold = 0.1;
    params.iouThreshold = 0.5;
    params.modelPath = "best.onnx";
    params.imgSize = {512, 512};
#ifdef USE_CUDA
    params.cudaEnable = true;

    // GPU FP32 inference
    params.modelType = YOLO_DETECT_V8;
    // GPU FP16 inference
    // Note: change fp16 onnx model
    // params.modelType = YOLO_DETECT_V8_HALF;

#else
    // CPU inference
    params.modelType = YOLO_DETECT_V8;
    params.cudaEnable = false;

#endif
    yoloDetector->CreateSession(params);
    // Detector(yoloDetector);
}

cv::Mat DetectTest(cv::Mat &img)
{
    cv::Mat detect_Img = videoDetector(yoloDetector, img);
    return detect_Img;
}

// void ClsTest()
// {
//     YOLO_V8* yoloDetector = new YOLO_V8;
//     std::string model_path = "best.onnx";
//     ReadCocoYaml(yoloDetector);
//     DL_INIT_PARAM params{ model_path, YOLO_CLS, {224, 224} };
//     yoloDetector->CreateSession(params);
//     Classifier(yoloDetector);
// }
